{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# How to define graph state\n",
        "\n",
        "This how to guide will cover different ways to define the state of your graph.\n",
        "\n",
        "## Prerequisites\n",
        "\n",
        "- [State conceptual guide](/langgraphjs/concepts/low_level/#state) - Conceptual guide on defining the state of your graph.\n",
        "- [Building graphs](/langgraphjs/tutorials/quickstart/) - This how-to assumes you have a basic understanding of how to build graphs.\n",
        "\n",
        "## Setup\n",
        "\n",
        "This guide requires installing the `@langchain/langgraph`, and `@langchain/core` packages:\n",
        "\n",
        "```bash\n",
        "npm install @langchain/langgraph @langchain/core\n",
        "```\n",
        "\n",
        "## Getting started\n",
        "\n",
        "The `Annotation` function is the recommended way to define your graph state for new `StateGraph` graphs. The `Annotation.Root` function is used to create the top-level state object, where each field represents a channel in the graph.\n",
        "\n",
        "Here's an example of how to define a simple graph state with one channel called `messages`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "import { BaseMessage } from \"@langchain/core/messages\";\n",
        "import { Annotation } from \"@langchain/langgraph\";\n",
        "\n",
        "const GraphAnnotation = Annotation.Root({\n",
        "  // Define a 'messages' channel to store an array of BaseMessage objects\n",
        "  messages: Annotation<BaseMessage[]>({\n",
        "    // Reducer function: Combines the current state with new messages\n",
        "    reducer: (currentState, updateValue) => currentState.concat(updateValue),\n",
        "    // Default function: Initialize the channel with an empty array\n",
        "    default: () => [],\n",
        "  })\n",
        "});"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Each channel can optionally have `reducer` and `default` functions:\n",
        "- The `reducer` function defines how new values are combined with the existing state.\n",
        "- The `default` function provides an initial value for the channel.\n",
        "\n",
        "For more information on reducers, see the [reducers conceptual guide](/langgraphjs/concepts/low_level/#reducers)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "const QuestionAnswerAnnotation = Annotation.Root({\n",
        "  question: Annotation<string>,\n",
        "  answer: Annotation<string>,\n",
        "});"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Above, all we're doing is defining the channels, and then passing the un-instantiated `Annotation` function as the value. It is important to note we always pass in the TypeScript type of each channel as the first generics argument to `Annotation`. Doing this ensures our graph state is type safe, and we can get the proper types when defining our nodes. Below shows how you can extract the typings from the `Annotation` function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "type QuestionAnswerAnnotationType = typeof QuestionAnswerAnnotation.State;"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This is equivalent to the following type:\n",
        "\n",
        "```typescript\n",
        "type QuestionAnswerAnnotationType = {\n",
        "  question: string;\n",
        "  answer: string;\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Merging states\n",
        "\n",
        "If you have two graph state annotations, you can merge the two into a single annotation by using the `spec` value:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "const MergedAnnotation = Annotation.Root({\n",
        "  ...QuestionAnswerAnnotation.spec,\n",
        "  ...GraphAnnotation.spec,\n",
        "})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The type of the merged annotation is the intersection of the two annotations:\n",
        "\n",
        "```typescript\n",
        "type MergedAnnotation = {\n",
        "  messages: BaseMessage[];\n",
        "  question: string;\n",
        "  answer: string;\n",
        "}\n",
        "```\n",
        "\n",
        "Finally, instantiating your graph using the annotations is as simple as passing the annotation to the `StateGraph` constructor:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [],
      "source": [
        "import { StateGraph } from \"@langchain/langgraph\";\n",
        "\n",
        "const workflow = new StateGraph(MergedAnnotation);"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## State channels\n",
        "\n",
        "The `Annotation` function is a convenience wrapper around the low level implementation of how states are defined in LangGraph. Defining state using the `channels` object (which is what `Annotation` is a wrapper of) is still possible, although not recommended for most cases. The below example shows how to implement a graph using this pattern:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "import { StateGraph } from \"@langchain/langgraph\";\n",
        "\n",
        "interface WorkflowChannelsState {\n",
        "  messages: BaseMessage[];\n",
        "  question: string;\n",
        "  answer: string;\n",
        "}\n",
        "\n",
        "const workflowWithChannels = new StateGraph<WorkflowChannelsState>({\n",
        "  channels: {\n",
        "    messages: {\n",
        "      reducer: (currentState, updateValue) => currentState.concat(updateValue),\n",
        "      default: () => [],\n",
        "    },\n",
        "    question: null,\n",
        "    answer: null,\n",
        "  }\n",
        "});"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Above, we set the value of `question` and `answer` to `null`, as it does not contain a default value. To set a default value, the channel should be implemented how the `messages` key is, with the `default` factory returning the default value. The `reducer` function is optional, and can be added to the channel object if needed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Using Zod\n",
        "\n",
        "If you want to add runtime validation to your state, you can use Zod instead of the `Annotation` function for state definition. You can also pass in your custom `reducer` and `default` factories as well by importing `@langchain/langgraph/zod`, which will extend Zod with LangGraph specific methods.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import \"@langchain/langgraph/zod\";\n",
        "import { z } from \"zod\";\n",
        "\n",
        "const AgentState = z.object({\n",
        "  messages: z\n",
        "    .array(z.string())\n",
        "    .default(() => [])\n",
        "    .langgraph.reducer(\n",
        "      (a, b) => a.concat(Array.isArray(b) ? b : [b]),\n",
        "      z.union([z.string(), z.array(z.string())])\n",
        "    ),\n",
        "  question: z.string(),\n",
        "  answer: z.string().min(1),\n",
        "});\n",
        "\n",
        "const graph = new StateGraph(AgentState);\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "TypeScript",
      "language": "typescript",
      "name": "tslab"
    },
    "language_info": {
      "codemirror_mode": {
        "mode": "typescript",
        "name": "javascript",
        "typescript": true
      },
      "file_extension": ".ts",
      "mimetype": "text/typescript",
      "name": "typescript",
      "version": "3.7.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}
